from fastapi import FastAPI
from pydantic import BaseModel
from langgraph.graph import StateGraph, START, END
from typing import TypedDict, Annotated
import asyncio

# 1. 定义 LangGraph 状态和节点
class WorkflowState(TypedDict):
    input: str  # 输入内容
    result: str  # 处理结果

def process_node(state: WorkflowState) -> WorkflowState:
    # 示例：简单处理（实际可替换为复杂工作流逻辑）
    return {"input": state["input"], "result": f"LangGraph 处理: {state['input'].upper()}"}

# 构建工作流
workflow = StateGraph(WorkflowState)
workflow.add_node("process", process_node)
workflow.add_edge(START, "process")
workflow.add_edge("process", END)
compiled_graph = workflow.compile()

# 2. 用 FastAPI 暴露接口
app = FastAPI()

class InputRequest(BaseModel):
    user_input: str

@app.post("/run-workflow")
async def run_workflow(request: InputRequest):
    # 调用 LangGraph 工作流
    result = await compiled_graph.ainvoke({"input": request.user_input})
    return {"status": "success", "result": result["result"]}

# 启动服务：uvicorn main:app --host 0.0.0.0 --port 8000